import cv2, os, time, datetime
from ocr_onnx import det_rec_functions
from PIL import Image, ImageDraw, ImageFont
import numpy as np


class CamaroCap(object):
    def __init__(self, ocr_sys):
        self.camNo = 0
        self.fpsCount = None
        self.startTime = None
        self.fontStyle = ImageFont.truetype("simsun.ttc", 30, encoding="utf-8")
        # OCR-检测-识别
        self.ocr_sys = ocr_sys  # det_rec_functions(None)

    def open(self):
        print('camNo: ' + str(self.camNo))
        self.cap = cv2.VideoCapture(int(self.camNo))
        camfps = self.cap.get(cv2.CAP_PROP_FPS)

        self.fpsCount = 0
        self.startTime = time.time()
        isRunning = False

        while self.cap.isOpened():
            isRunning = True
            retval, frame = self.cap.read()
            # print("retval:" + str(retval))
            # print("frame:" + str(frame))

            if frame is not None:
                fps = self.calcFPS(frame)
                frame = self.cv2_chinese_text(frame, 'FPS:' + str(fps), 0, 0)
                if self.ocr_sys is not None:
                    rec = self.ocr(frame)
                    if rec is None or len(rec) == 0:
                        rec = '未识别'
                    frame = self.cv2_chinese_text(frame, 'REC:' + str(rec), 0, 50)

                cv2.imshow('frame', frame)
            # 键盘输入空格暂停，输入q退出
            key = cv2.waitKey(1) & 0xff
            if key == ord(" "):
                cv2.waitKey(0)
            if key == ord("q"):
                break
            # time.sleep(1 / camfps if camfps > 0 else 60)  # 按原帧率播放
        if not isRunning:
            print('cap is not Opened.')

    def setText(self, frame, text, x, y):
        # 不支持中文
        cv2.putText(frame,
                    text, (x, y),
                    cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255),
                    3)
        return frame

    def cv2_chinese_text(self, img, text, x, y,
                         textColor=(255, 0, 0), textSize=30):
        if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
            img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        # 创建一个可以在给定图像上绘图的对象
        draw = ImageDraw.Draw(img)
        # 绘制文本
        draw.text((x, y), text, textColor, font=self.fontStyle)
        # 转换回OpenCV格式
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

    def calcFPS(self, frame):
        self.fpsCount += 1
        fps = 0
        if (time.time() - self.startTime) != 0:  # 实时显示帧数
            fps = float('%.1f' % (self.fpsCount / (time.time() - self.startTime)))
            print('fps:' + str(fps))
            # self.setText(frame, str(fps), 100, 100)
            self.fpsCount = 0
            self.startTime = time.time()
        return fps

    def ocr(self, frame):
        self.ocr_sys.img = frame.copy()
        # 得到检测框
        dt_boxes = self.ocr_sys.get_boxes()
        # 识别 results: 单纯的识别结果，results_info: 识别结果+置信度
        results, results_info = self.ocr_sys.recognition_img(dt_boxes)
        print('result: ' + str(results))
        # self.setText(frame, str(results), 100, 200)
        return results

    def release(self):
        self.cap.release()
        cv2.destroyAllWindows()


if __name__ == '__main__':
    cam = CamaroCap(det_rec_functions(None))
    cam.open()
    cam.release()
